import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['KaiTi']
plt.rcParams['axes.unicode_minus'] = False
import seaborn as sns
import pandas as pd
import numpy as np
from collections import defaultdict
from sklearn.manifold import TSNE
from sklearn.decomposition import PCA
from scipy.optimize import linear_sum_assignment
from matplotlib.colors import ListedColormap
import io
import base64
def custom_confusion_matrix(y_true, y_pred, true_classes, pred_classes):
    confusion_dict = {}

    for true_label, pred_label in zip(y_true, y_pred):
        # Ensuring the extraction of scalar values if the inputs are numpy arrays
        true_label = true_label.item() if isinstance(true_label, np.ndarray) else true_label
        pred_label = pred_label.item() if isinstance(pred_label, np.ndarray) else pred_label

        # Ensure the classes are within expected categories
        if true_label in true_classes and pred_label in pred_classes:
            key = (true_label, pred_label)
            confusion_dict[key] = confusion_dict.get(key, 0) + 1

    return confusion_dict

def custom_sort_key(series, kwn, unkwn):
    # 假设 kwn 和 unkwn 是你定义好的、代表已知和未知标签顺序的列表
    all_labels = list(set(series))  # 获取 series 中所有独特的标签
    custom_order = {label: idx for idx, label in enumerate(kwn + unkwn)}  # 构建映射
    # 使用映射来生成排序依据的序列
    return series.map(custom_order)

def draw_output_pic(output, predict_labels, raw_label, reconstruct_threshold, kwn, unkwn):
    # plt.figure()

    output_sorted = output.sort_values('raw label', key=lambda x: custom_sort_key(x, kwn, unkwn))
    sns.jointplot(x='raw label', y='loss', data=output_sorted)
    plt.axhline(y=reconstruct_threshold, color='red', linewidth=2, linestyle='-', )
    plt.text(0.1, reconstruct_threshold, 'Threshold', color='red', fontsize=12, ha='center', va='bottom')
    # plt.show()

    confusion_dict = custom_confusion_matrix(
        [item.item() if isinstance(item, np.ndarray) else item for item in raw_label],
        [item.item() if isinstance(item, np.ndarray) else item for item in predict_labels],
        kwn + unkwn,
        kwn + ['-1']
    )

    df_data = defaultdict(dict)
    for (x, y), value in confusion_dict.items():
        df_data[x][y] = value
    df = pd.DataFrame(df_data)
    # 归一化处理
    total_counts = df.sum(axis=0)  # 每个类别的总样本数
    df_normalized = df.divide(total_counts, axis=1) * 100  # 归一化并转换为百分比

    valid_indices = [idx for idx in kwn + ['-1'] if idx in df.index]
    # 按行索引排序
    if valid_indices:  # 确保列表不为空
        df_sorted_rows = df_normalized.reindex(valid_indices)
    else:
        print("No valid indices found for sorting.")
    valid_col = [idx for idx in kwn + unkwn if idx in df.columns]
    # 按列名排序
    df_sorted_both = df_sorted_rows[valid_col]
    # 绘制热力图
    plt.figure(figsize=(8, 6))
    sns.heatmap(df_sorted_both, annot=True, fmt=".2f", cmap='Blues', linewidths=.5, cbar=False,
                xticklabels=valid_col, yticklabels=valid_indices)  #[::-1]
    plt.title("Confusion Matrix", fontsize=16)
    plt.ylabel("Predicted Label", fontsize=14)
    plt.xlabel("True Label", fontsize=14)
    return df_sorted_both, valid_col, valid_indices

def draw_confuse_pic(predict_labels, raw_label, kwn, unkwn):

    confusion_dict = custom_confusion_matrix(
        [item.item() if isinstance(item, np.ndarray) else item for item in raw_label],
        [item.item() if isinstance(item, np.ndarray) else item for item in predict_labels],
        kwn + unkwn,
        kwn + ['-1']
    )

    df_data = defaultdict(dict)
    for (x, y), value in confusion_dict.items():
        df_data[x][y] = value
    df = pd.DataFrame(df_data)
    # 归一化处理
    total_counts = df.sum(axis=0)  # 每个类别的总样本数
    df_normalized = df.divide(total_counts, axis=1) * 100  # 归一化并转换为百分比

    valid_indices = [idx for idx in kwn + ['-1'] if idx in df.index]
    # 按行索引排序
    if valid_indices:  # 确保列表不为空
        df_sorted_rows = df.reindex(valid_indices)
    else:
        print("No valid indices found for sorting.")
    valid_col = [idx for idx in kwn + unkwn if idx in df.columns]
    # 按列名排序
    df_sorted_both = df_sorted_rows[valid_col]
    # 将 NaN 值替换为 0
    df_sorted_both.fillna(0, inplace=True)
    df_sorted_both = df_sorted_both.astype('int')
    # 创建热图
    plt.figure(figsize=(8, 6))
    sns.heatmap(df_sorted_both, annot=True, fmt="d", cmap='viridis', linewidths=.5, cbar=True,
                xticklabels=valid_col, yticklabels=valid_indices) #[::-1]
    plt.title("混淆矩阵", fontsize=16)
    plt.ylabel("预测标签", fontsize=14)
    plt.xlabel("实际标签", fontsize=14)
    # 将图像保存到字节流
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)

    # 将字节流编码为Base64
    image = base64.b64encode(buf.getvalue()).decode('utf-8')

    return df_sorted_both, valid_col, valid_indices, image


def draw_pca(features, predict_labels):
    # 确保 predict_labels 是 NumPy 数组以便进行索引
    predict_labels = np.array(predict_labels)

    # 初始化 PCA 模型
    pca = PCA(n_components=2)  # n_components 设置为 2，即降维到二维
    # 对数据进行降维
    X_pca = pca.fit_transform(features)

    plt.figure()
    unique_labels = np.unique(predict_labels)
    for label in unique_labels:
        indices = predict_labels == label
        plt.scatter(X_pca[indices, 0], X_pca[indices, 1],
                    label=f'Class {label}')

    plt.title('PCA Visualization with Class Labels')
    plt.xlabel('Principal Component 1')
    plt.ylabel('Principal Component 2')
    plt.legend()
    # plt.show()

    return X_pca, pca

def draw_tsne(features, predict_labels):

    # 假设labels是与features_train对应的一维数组，包含了样本的类别标签
    unique_labels = np.unique(predict_labels)
    # 初始化PCA并将其限制在前两或三个主成分
    X = features
    y = np.array(predict_labels)
    # 初始化t-SNE模型
    tsne = TSNE(n_components=2, random_state=42)  # n_components设置为2，即降维到二维
    # 对数据进行降维
    X_tsne = tsne.fit_transform(X)
    plt.figure()
    for label in unique_labels:
        plt.scatter(X_tsne[y == label, 0], X_tsne[y == label, 1],
                    label=f'Class {label}')
    plt.title('t-SNE Visualization with Class Labels')
    plt.xlabel('Principal Component 1')
    plt.ylabel('Principal Component 2')
    plt.legend()

def draw_load_pca(features, predict_labels, pca, centers):
    # 确保 predict_labels 是 NumPy 数组以便进行索引
    predict_labels = np.array(predict_labels)

    # 对数据进行降维
    X_pca = pca.transform(features)

    plt.figure()
    unique_labels = np.unique(predict_labels)
    for label in unique_labels:
        indices = predict_labels == label
        plt.scatter(X_pca[indices, 0], X_pca[indices, 1],
                    label=f'Class {label}')
    plt.scatter(centers[:, 0], centers[:, 1])

    plt.title('PCA Visualization with Class Labels')
    plt.xlabel('Principal Component 1')
    plt.ylabel('Principal Component 2')
    plt.legend()
    # plt.show()

    return X_pca

    # 映射聚类标签到真实标签


# 映射聚类标签到真实标签
def map_clusters_to_labels(cluster_labels, true_labels, unkwn):

    true_labels = np.array(true_labels)
    true_label_type = np.unique(true_labels)
    num_classes = len(np.unique(true_labels))
    pre_classes = len(np.unique(cluster_labels))
    cost_matrix = np.zeros((num_classes, num_classes))
    for i in range(pre_classes):
        for j in range(num_classes):
            label = true_label_type[j]
            cost_matrix[i, j] = -np.sum((cluster_labels == i) & (true_labels == label))
    row_ind, col_ind = linear_sum_assignment(cost_matrix)

    # 将聚类标签映射为字符串类型
    mapping = {i: true_label_type[col] for i, col in enumerate(col_ind)}
    mapped_labels = np.array([mapping.get(label, label) for label in cluster_labels], dtype=object)
    return mapped_labels
